#include <stdio.h>
#include "serverControl.h"
#include <time.h>
#include <semaphore.h>
#define MAX_THREADS 1
#if _WIN32

#include <winsock2.h>
#include <windows.h>
#include <process.h>


#pragma comment(lib, "ws2_32.lib")
//互斥量
HANDLE cacheLock; //保护dynamicCache
HANDLE tableLock; //保护id映射表
#elif __linux__

#include <sys/socket.h>
#include <netinet/in.h>
#include <pthread.h>
#include <unistd.h>
pthread_mutex_t cacheLock=PTHREAD_MUTEX_INITIALIZER;
pthread_mutex_t tableLock=PTHREAD_MUTEX_INITIALIZER;
#endif
//创建线程时传入的参数

struct threadVariable {
#if _WIN32
    SOCKET s;
#elif __linux__
    int s;
#endif
    struct sockaddr_in *si_in;
    int len;
    struct hashmap_s *staticHashmap;
    struct hashmap_s *dynamicHashmap;
    char DNSserver[16];
    char *packet;
};
sem_t threadSemaphore; // 信号量,监控同时建立的最大线程数量
#if _WIN32
DWORD WINAPI ProcessPacket(LPVOID lpParam) {
    struct threadVariable *threadVariable = (struct threadVariable *) lpParam;
    SOCKET s = threadVariable->s;
    struct sockaddr_in *si_in = threadVariable->si_in;
    int len = threadVariable->len;
    struct hashmap_s *staticHashmap = threadVariable->staticHashmap;
    struct hashmap_s *dynamicHashmap = threadVariable->dynamicHashmap;
    char *DNSserver = threadVariable->DNSserver;
    char *packet = threadVariable->packet;
    processQuery(s, si_in, staticHashmap, dynamicHashmap, packet, len, DNSserver);
    free(threadVariable);
    free(packet);
    sem_post(&threadSemaphore);
    return 0;
}

#elif __linux__
void* ProcessPacketForUnix(void *lpParam){
    struct threadVariable *threadVariable = (struct threadVariable *) lpParam;
    int s = threadVariable->s;
    struct sockaddr_in *si_in = threadVariable->si_in;
    int len = threadVariable->len;
    struct hashmap_s *staticHashmap = threadVariable->staticHashmap;
    struct hashmap_s *dynamicHashmap = threadVariable->dynamicHashmap;
    char *DNSserver = threadVariable->DNSserver;
    char *packet = threadVariable->packet;
    processQuery(s, si_in, staticHashmap, dynamicHashmap, packet, len, DNSserver);
    free(threadVariable);
    free(packet);
    sem_post(&threadSemaphore);
    return 0;
}
#endif

int main() {
    //初始化cache,对应用户设置的静态规则，从ISP获取的动态规则，随机ID
    struct hashmap_s staticCache;
    struct hashmap_s dynamicCache;
    printf("+-+-+-+-+-+-+-+-+-+\n"
           "|D|N|S|-|R|E|L|A|Y|\n"
           "+-+-+-+-+-+-+-+-+-+\n");
    //get local data and time
    time_t now;
    time(&now);
    struct tm *local = localtime(&now);
    printf("Build by Ye Wenyu at %s\n", asctime(local));
    printf("Thanks to the 'hashmap.h' open source project on Github for providing the hashmap implementation!\n");
    printf("https://github.com/sheredom/hashmap.h\n");
    char *DNSserver = malloc(16);
    printf("Select your own DNS server address[Y/N=Cloudflare DNS]:");
    char choice;
    scanf("%c", &choice);
    if (choice == 'Y' || choice == 'y') {
        printf("\nInput your DNS server address:");
        scanf("%s", DNSserver);
    } else {
        strcpy(DNSserver, "1.1.1.1");
    }
//        strcpy(DNSserver, "1.1.1.1");
    sem_init(&threadSemaphore, 0, MAX_THREADS);
#if _WIN32
    //初始化互斥量
    cacheLock = CreateMutex(NULL, FALSE, NULL);
    tableLock = CreateMutex(NULL, FALSE, NULL);
#elif __linux__
    pthread_mutex_init(&cacheLock, NULL);
    pthread_mutex_init(&tableLock, NULL);
#endif
    initServer(&staticCache, "..//test.txt", &dynamicCache);
    char inBuffer[1024];
    printf("awake dynamicCache refresh thread\n");
#if _WIN32
    sleep(1);
    _beginthread(refreshCacheForWin32, 0, &dynamicCache);
    int value;
    sem_getvalue(&threadSemaphore, &value);
    printf("awake threadSemaphore value: %d\n", value);
    SOCKET s;
    struct sockaddr_in server, si_other;
    int slen;
    WSADATA wsa;

    slen = sizeof(si_other);

    //Initialise winsock
    printf("\nInitialising Winsock...");
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0) {
        printf("Failed. Error Code : %d", WSAGetLastError());
        exit(EXIT_FAILURE);
    }
    printf("Initialised.\n");

    //Create a socket
    if ((s = socket(AF_INET, SOCK_DGRAM, 0)) == INVALID_SOCKET) {
        printf("Could not create socket : %d", WSAGetLastError());
    }
    printf("Socket created.\n");
#elif __linux__
    sleep(1);
    pthread_t cahceThread;
    pthread_create(&cahceThread, NULL, refreshCacheForUnix, &dynamicCache);
    int s;
    struct sockaddr_in server, si_other;
    int slen;
    slen = sizeof(si_other);
    //Create a socket
    if ((s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1) {
        printf("Could not create socket");
    }
    printf("Socket created\n");
#endif
    //Prepare the sockaddr_in structure
    server.sin_family = AF_INET;
    server.sin_addr.s_addr = inet_addr("127.0.0.1");
#if _WIN32
    server.sin_port = htons(53);
#elif __linux__
    server.sin_port = htons(8000);
#endif
    printf("server ip:%s\n", inet_ntoa(server.sin_addr));
#if _WIN32
    //Bind
    if (bind(s, (struct sockaddr *) &server, sizeof(server)) == SOCKET_ERROR) {
        printf("Bind failed with error code : %d", WSAGetLastError());
        exit(EXIT_FAILURE);
    }
    puts("Bind port 53");
#elif __linux__
    //Bind
    if (bind(s, (struct sockaddr *) &server, sizeof(server)) == -1) {
        printf("Bind failed\n");
        exit(EXIT_FAILURE);
    }
    puts("Bind port 8000");
#endif
    //keep listening for data
    while (1) {
#if _WIN32
        //等待计数器大于0

        sem_getvalue(&threadSemaphore, &value);
//        if(value==0)
//            sleep(1);
        sem_wait(&threadSemaphore);
        int len = recvfrom(s, inBuffer, sizeof(inBuffer), 0, (struct sockaddr *) &si_other,
                           &slen);
#elif __linux__
        int len = recvfrom(s, inBuffer, sizeof(inBuffer), 0, (struct sockaddr *) &si_other,
                           (socklen_t *) &slen);
#endif
        if (len < 0) {
            printf("Waiting!\n");
            //接收失败需要恢复计数
            sem_post(&threadSemaphore);
//            sleep(1);
            continue;
        }
#if _WIN32
        HANDLE hThread;
        DWORD threadId;
#elif __linux__
        pthread_t hThread;

#endif
        //malloc threadVariable在ProcessPacket中释放

        struct threadVariable *threadVariable = (struct threadVariable *) malloc(sizeof(struct threadVariable));
        threadVariable->s = s;
        threadVariable->si_in = &si_other;
        threadVariable->len = len;
        threadVariable->staticHashmap = &staticCache;
        threadVariable->dynamicHashmap = &dynamicCache;
        threadVariable->packet = (char *) malloc(len);
        memcpy(threadVariable->packet, inBuffer, len); // 复制数据到独立的缓冲区
        strcpy(threadVariable->DNSserver, DNSserver);

#if _WIN32
        hThread = CreateThread(NULL, 0, ProcessPacket, (LPVOID) threadVariable, 0, &threadId);
        if (hThread == NULL) {
            printf("CreateThread failed\n");
            continue;
        }
        CloseHandle(hThread);
#elif __linux__

        if(0!=pthread_create(&hThread, NULL, ProcessPacketForUnix, (void *) threadVariable)){
            printf("pthread_create failed\n");
            continue;
        }
        pthread_detach(hThread);
#endif
        threadVariable = NULL;
    }
    sem_destroy(&threadSemaphore);
#if _WIN32
    closesocket(s);
    WSACleanup();
#elif __linux__
    close(s);
#endif
    DNSFree();
    return 0;
}

